
import chisel3._
// import chisel3.util._
import chiseltest._
import org.scalatest.{Matchers, FlatSpec}
import scala.util.Random

import comlib._
import comlib.arbiter._

class SelectUnitSpec extends FlatSpec with ChiselScalatestTester with Matchers {

  val rnd = new Random

  it should "支持无请求，无应答，屏蔽透传" in {
    test( new SelectUnit(1) ){ dut =>
      for( i <- 0 until 10 ){
        dut.io.r(0).poke( false.B )
        dut.io.g(0).expect( false.B )

        val ei = rnd.nextInt(2)==0
        dut.io.ei(0).poke( ei.B )
        dut.io.wo(0).expect( ei.B )

        val si = rnd.nextInt(2)==0
        dut.io.si(0).poke( si.B )
        dut.io.no(0).expect( si.B )

        val wi = rnd.nextInt(2)==0
        dut.io.wi(0).poke( wi.B )
        dut.io.eo(0).expect( wi.B )

        val ni = rnd.nextInt(2)==0
        dut.io.ni(0).poke( ni.B )
        dut.io.so(0).expect( ni.B )
        
      }
    }
  }

}

class RoundRobinSpec extends FlatSpec with ChiselScalatestTester with Matchers {

  val rnd = new Random

  it should "支持简单1D SP调度" in {

    // 定义模块
    class MyModule extends Module{
      val io = IO(new Bundle{
        val req = Input(UInt(4.W))
        val gnt = Output(UInt(4.W))
      })
      val sp = Arbiter.sp(VecInit(io.req.asBools))
      io.gnt := sp.io.gnt.asUInt()
    }

    // 定义激励与预期
    val res0 = Seq(
      "b1111".U -> "b0001".U,
      "b1110".U -> "b0010".U,
      "b1100".U -> "b0100".U,
      "b1000".U -> "b1000".U,
      "b0000".U -> "b0000".U,
      "b1010".U -> "b0010".U,
      "b0101".U -> "b0001".U,
    )

    // 测试
    test( new MyModule ){ dut =>
      for( (r, g) <- res0 ){
        dut.io.req.poke(r)
        dut.io.gnt.expect(g)
      }
    }

  }


  it should "支持SP-RR调度" in {

    class MyModule extends Module{
      val io = IO(new Bundle{
        val req = Input( Vec(4, Vec(4, Bool())) )
        val gnt = Output( Vec(4, Vec(4, Bool())) )
        val ptr = Input( Vec(4, Bool()) )
      })

      val sprr = Module(new Arbiter( 4, 4, 1, colMode=1 ))
      sprr.io.rmsk := sprr.rmsk(false.B)

      sprr.io.r := sprr.req( io.req )
      sprr.io.cmsk := sprr.cptr( io.ptr )
      io.gnt := sprr.io.gnt

    }

    test( new MyModule ){ dut =>

      // ==== 随机分配测试
      val cptr3 = rnd.nextInt(4)   // 随机指针
      val in_num = rnd.nextInt(4)
      val in_rows = Random.shuffle( (0 until in_num).toList ).toList.sorted

      val cptr_seq = (0 until 4).map { case i => if(i<cptr3) true.B else false.B }   // binary To temp

      // Request
      for( r <- 0 until 4; c <- 0 until 4 )
        if( in_rows.contains(r)  ) dut.io.req(r)(c).poke( true.B )
        else dut.io.req(r)(c).poke( false.B )

      // Pointer
      (0 until 4).foreach { case i => dut.io.ptr(i).poke(cptr_seq(i)) }

      dut.clock.step()

      // Expect
      for( r <- 0 until 4; c <- 0 until 4 ){
        in_rows.indexOf( r ) match{
          case x if( x <0 ) =>
            dut.io.gnt(c)(r).expect( false.B, f"Point=${cptr3} and inRows=${in_rows}" )
          case ofst =>
            val cpos = (cptr3 + ofst) % 4
            if( c == cpos ) dut.io.gnt(c)(r).expect( true.B, f"Point=${cptr3} and inRows=${in_rows}" )
            else dut.io.gnt(c)(r).expect( false.B, f"Point=${cptr3} and inRows=${in_rows}" )
        }
      }
      
    }

  }

  it should "支持SP-RR调度，解决挤气泡问题" in {

    class MyModule extends Module{
      val io = IO(new Bundle{
        val req = Input( Vec(4, Vec(4, Bool())) )
        val gnt = Output( Vec(4, Vec(4, Bool())) )
        val cmo = Output( UInt(7.W) )
      })

      val ptr = RegInit(VecInit(Seq.fill(4)(false.B)))    // 指针
      val sprr = Module(new Arbiter( 4, 4, 1, colMode=1 ))    // 调度器
      sprr.io.rmsk := sprr.rmsk(false.B)    // 默认屏蔽

      sprr.io.r := sprr.req( io.req )   // 请求赋值
      sprr.io.cmsk := sprr.cptr(ptr)   // 指针赋值
      io.gnt := sprr.io.gnt    // Grant输出

      val cmo = sprr.io.cmo(0)
      ptr := Mux( cmo(3), VecInit(cmo.slice(4,7):+false.B), VecInit(cmo.slice(0,4)) )   // 指针更新
      io.cmo := sprr.io.cmo(0).asUInt()

      // chisel3.printf(p"${sprr.show()}")
      
    }

    test( new MyModule ){ dut =>

      var ptrb = 0

      for( it <- 0 until 10 ){

        val in_num = rnd.nextInt(4)   // 随机输入数据个数
        val in_rows = Random.shuffle( (0 until in_num).toList ).toList.sorted   // 随机输入数据通道号

        // Request
        for( r <- 0 until 4; c <- 0 until 4 )
          if( in_rows.contains(r)  ) dut.io.req(r)(c).poke( true.B )
          else dut.io.req(r)(c).poke( false.B )

        // Expect
        for( r <- 0 until 4; c <- 0 until 4 ){
          in_rows.indexOf( r ) match{
            case x if( x <0 ) =>
              dut.io.gnt(c)(r).expect( false.B, f"Point=${ptrb} and inRows=${in_rows} it=${it}" )
            case ofst =>
              val cpos = (ptrb + ofst) % 4
              if( c == cpos ) dut.io.gnt(c)(r).expect( true.B, f"Point=${ptrb} and inRows=${in_rows} it=${it}" )
              else dut.io.gnt(c)(r).expect( false.B, f"Point=${ptrb} and inRows=${in_rows} it=${it}" )
          }
        }

        // Update Pointer
        ptrb = (ptrb+in_num) % 4

        dut.clock.step()

      }

    }

  }

  
  it should "支持rmsk信息扩展" in {

    class MyModule(ncol:Int, nrow:Int, npri:Int) extends Module{
      val io = IO(new Bundle{
        val s = Input(Bool())
        val res_s = Output(Vec(npri, Vec(nrow, Bool())))
      })

      val arb = Module(new Arbiter( ncol, nrow, npri ))    // 调度器
      arb.io.rmsk := arb.rmsk( false.B )
      arb.io.cmsk := arb.cmsk( false.B )
      arb.io.r := arb.req( false.B )

      io.res_s := arb.rmsk( io.s )      
    }

    val rnd = new Random
    val npri = rnd.nextInt(4) + 1
    val nrow = rnd.nextInt(4) + 1
    val ncol = rnd.nextInt(4) + 1

    Predef.printf(f"\nCurrent npri=${npri} nrow=${nrow} ncol=${ncol} \n")

    test( new MyModule(ncol, nrow, npri) ){ dut =>

      val s = if( rnd.nextInt(1)==0 ) false.B else true.B 
      dut.io.s.poke( s )
      for( p <- 0 until npri; r <- 0 until nrow ){
        dut.io.res_s(p)(r).expect( s )
      }

    }


    test( new Module{
      val io = IO(new Bundle{
        val rmsk_all = Input(Bool())
        val rmsk_vec = Input(Vec(4, Bool()))
        val gnt = Output(Vec(4, Vec(4, Bool())))
      })

      val rmsk = Wire(Vec(1, Vec(4, Bool())))
      val arb = Module( new Arbiter( 4, 4, 1 ) )
      arb.io.r := arb.req( true.B )
      val all = arb.rmsk( io.rmsk_all )
      val vec = arb.rmsk( io.rmsk_vec )
      for( p <- 0 until 1; r <- 0 until 4 ) rmsk(p)(r) := all(p)(r) | vec(p)(r)
      arb.io.rmsk := rmsk
      arb.io.cmsk := arb.cmsk( false.B )
      io.gnt := arb.io.gr(0)

    } ){dut =>

      dut.io.rmsk_all.poke( true.B )
      dut.io.rmsk_vec.foreach( _.poke( false.B ) )
      dut.io.gnt.foreach( _.foreach( _.expect( false.B ) ) )

      dut.io.rmsk_all.poke( false.B )
      dut.io.rmsk_vec.foreach( _.poke( true.B ) )
      dut.io.gnt.foreach( _.foreach( _.expect( false.B ) ) )

    }


  }

  it should "支持until到0选择策略" in {

    class MyModule extends Module{
      val io = IO(new Bundle{
        val req = Input(UInt(4.W))
        val ptr = Input(UInt(4.W))
        val gnt = Output(UInt())
      })

      val arb = Module(new Arbiter( 4, colMode=1, selMode=3 ))    // 调度器
      arb.io.rmsk := arb.rmsk( false.B )
      arb.io.cmsk := arb.cptr( VecInit(io.ptr.asBools()) )
      arb.io.r := arb.req( VecInit(io.req.asBools()) )
      io.gnt := arb.io.gnt.asUInt()

    }

    test( new MyModule ){ dut =>

      val result = Seq(
        ("b1111".U, "b0000".U, "b1111".U),
        ("b0000".U, "b0000".U, "b0000".U),
        ("b0001".U, "b0000".U, "b0001".U),
        ("b0011".U, "b0000".U, "b0011".U),
        ("b0111".U, "b0000".U, "b0111".U),
        ("b1011".U, "b0000".U, "b0011".U),
        ("b1101".U, "b0000".U, "b0001".U),
        ("b1011".U, "b0000".U, "b0011".U),

        ("b1111".U, "b0001".U, "b1111".U),
        ("b1011".U, "b0001".U, "b0010".U),
        ("b1011".U, "b0011".U, "b0000".U),
        ("b1110".U, "b0011".U, "b1100".U),
      )


      for( (r,p,g) <- result ){
        dut.io.ptr.poke(p)
        dut.io.req.poke(r)
        dut.io.gnt.expect(g)
      }

    }
  }
 
}


class ArbiterSpec extends FlatSpec with ChiselScalatestTester with Matchers {

  it should "Support Squeeze Bubble Arbiter" in {

    test( new Module{
      val io = IO( new Bundle{
        val req = Input(UInt(4.W))
        val gnt = Output(Vec(4, UInt()))
      } )
      
      val arb = Arbiter.squeezeBubble( VecInit(io.req.asBools()) )
      io.gnt := VecInit(arb.io.gc(0).map( _.asUInt() ))

      // printf(p"${arb.show()}")

    } ){ dut=>

      val result = Seq(
        "b0000".U -> Seq("b0000".U, "b0000".U, "b0000".U, "b0000".U),
        "b0001".U -> Seq("b0001".U, "b0000".U, "b0000".U, "b0000".U),
      )

      for( (r, g) <- result ){

        dut.io.req.poke(r)
        g.zipWithIndex.foreach { case (gi,i) => dut.io.gnt(i).expect(gi) }
        dut.clock.step()

      }

    }
  }

  it should "Support 4 nodes Round-Robin Grant config" in {

    test( new Module{
      val io = IO( new Bundle{
        val gntr = Output(Vec(4, Vec(4, UInt())))
        val gnt = Output(Vec(4, Vec(4, UInt())))
      } )
      
      io.gntr := Arbiter.rrSelect(4, true)
      io.gnt := Arbiter.rrSelect(4, false)

    } ){ dut=>

      val resultr = Seq(
        Seq( 0.U, 1.U, 2.U, 3.U ),
        Seq( 3.U, 0.U, 1.U, 2.U ),
        Seq( 2.U, 3.U, 0.U, 1.U ),
        Seq( 1.U, 2.U, 3.U, 0.U ),
      )

      val result = Seq(
        Seq( 0.U, 1.U, 2.U, 3.U ),
        Seq( 1.U, 2.U, 3.U, 0.U ),
        Seq( 2.U, 3.U, 0.U, 1.U ),
        Seq( 3.U, 0.U, 1.U, 2.U ),
      )

      for( s <- 0 until 4; o <- 0 until 4 ){
        dut.io.gntr(s)(o).expect( resultr(s)(o) )
        dut.io.gnt(s)(o).expect( result(s)(o) )
      }

    }
  }

  it should "support roundRobin1 function" in {

    class MyModule(reverse:Boolean) extends Module{
      val io = IO( new Bundle{
        val in = Input(Vec(4, Bool()))
        val ptr = Input(UInt(2.W))
        val out = Output(Vec(4, Bool()))
      } )
      io.out := Arbiter.roundRobinD1( io.in, io.ptr, reverse )
    }

    test( new MyModule(false) ){ dut =>
      val testcast1 = Seq(
        (Seq(true.B, false.B, false.B, true.B), 0.U, Seq(true.B, false.B, false.B, true.B)),
        (Seq(true.B, false.B, false.B, true.B), 1.U, Seq(false.B, false.B, true.B, true.B)),
        (Seq(true.B, false.B, false.B, true.B), 2.U, Seq(false.B, true.B, true.B, false.B)),
        (Seq(true.B, false.B, false.B, true.B), 3.U, Seq(true.B, true.B, false.B, false.B)),
      )

      for( (in, ptr, out) <- testcast1 ){
        dut.io.in.zip(in).foreach {case (io,input) => io.poke(input)}
        dut.io.ptr.poke(ptr)
        dut.io.out.zip(out).foreach {case (io,output) => io.expect(output)}
      }
    }

    test( new MyModule(true) ){ dut =>
      val testcast1 = Seq(
        (Seq(true.B, false.B, false.B, true.B), 0.U, Seq(true.B, false.B, false.B, true.B)),
        (Seq(true.B, false.B, false.B, true.B), 1.U, Seq(true.B, true.B, false.B, false.B)),
        (Seq(true.B, false.B, false.B, true.B), 2.U, Seq(false.B, true.B, true.B, false.B)),
        (Seq(true.B, false.B, false.B, true.B), 3.U, Seq(false.B, false.B, true.B, true.B)),
      )

      for( (in, ptr, out) <- testcast1 ){
        dut.io.in.zip(in).foreach {case (io,input) => io.poke(input)}
        dut.io.ptr.poke(ptr)
        dut.io.out.zip(out).foreach {case (io,output) => io.expect(output)}
      }
    }
  }

}


